#encoding=utf-8
import cv2
from vision.core.detect import create_mtcnn_net, MtcnnDetector
import glob,os


if __name__ == '__main__':

    eval_net = 'onet' #'pnet' or 'rnet' or 'onet
    is_show = False # detection results visualization
    fddb_root = '/media/handsome/data/hanson/FDDB'

    pnet, rnet, onet = create_mtcnn_net(p_model_path="/media/handsome/data/hanson/WIDER_train/model_store/heils_wider/pnet/pnet_epoch_35_0.958786.pt",
                                        r_model_path="/media/handsome/data/hanson/WIDER_train/model_store/heils_wider/rnet/rnet_epoch_27_0.974149.pt",
                                        o_model_path="/media/handsome/data/hanson/WIDER_train/model_store/heils_wider/onet/onet_epoch_6_0.973742_0.0003.pt",
                                        use_cuda=True)
    if eval_net == 'pnet':
        mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=None, onet=None, min_face_size=24)
    elif eval_net == 'rnet':
        mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, onet=None, min_face_size=24)
    elif eval_net == 'onet':
        mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, onet=onet, min_face_size=24)


    ori_images_folder = os.path.join(fddb_root, 'originalPics')
    img_list = os.path.join(fddb_root, 'imList.txt')
    results = os.path.join(fddb_root, 'results.txt')
    img_paths = []
    with open(img_list, 'r')  as f:
        img_paths = f.readlines()

    f = open(results,'w')

    total_face = 0
    for idx, path in enumerate(img_paths):

        img_path = os.path.join(ori_images_folder, path.strip() + '.jpg')
        img = cv2.imread(img_path)
        img_bg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        bboxs, landmarks = mtcnn_detector.detect_face(img_bg)
        # print box_align

        face_num = bboxs.shape[0]
        total_face += face_num
        print('image id:{}, face num:{}, total num:{}, image path:{}'.format(idx, face_num, total_face, path))
        f.write(path)
        f.write('%d\n' % face_num)

        for i in range(face_num):
            x1 = int(bboxs[i][0])
            y1 = int(bboxs[i][1])
            x2 = int(bboxs[i][2])
            y2 = int(bboxs[i][3])
            w = x2 - x1
            h = y2 - y1
            score = bboxs[i][4]

            f.write('%d %d %d %d %f\n'%(x1, y1, w ,h ,score))

            cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2)
            if eval_net == 'onet':
                for j in range(5):
                    cv2.circle(img, (int(landmarks[i][2 * j]), int(landmarks[i][2 * j + 1])), 1, (255, 255, 0))

        if is_show:
            cv2.imshow("img", img)
            cv2.waitKey(0)

    f.close()